import time

import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

from scipy.stats import kurtosis


# 设置中文字体，解决显示中文乱码问题
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用 SimHei 字体，或者尝试 Microsoft YaHei
matplotlib.rcParams['axes.unicode_minus'] = False  # 防止负号显示问题

def check_orthogonality_and_determinant(matrix):
    n = matrix.size(0)
    identity = torch.eye(n, dtype=matrix.dtype, device=matrix.device)
    product = torch.matmul(matrix, matrix.t())  # 计算矩阵乘以其转置
    error = torch.norm(product - identity)  # Frobenius 范数计算误差
    matrix = matrix.to(torch.float32)
    determinant = torch.linalg.det(matrix)  # 计算行列式
    return error.item(), determinant.item()


def show(mp):
    plt.show()


def compute_global_and_channel_kurtosis_numpy(tensor):
    # 转换为 NumPy 数组
    tensor_np = tensor.detach().cpu().numpy()

    # 1. 计算总体峰度
    global_kurtosis = kurtosis(tensor_np.ravel())

    # 2. 合并前两个维度
    reshaped_tensor = tensor_np.reshape(-1, tensor_np.shape[-1])  # (14, hidden_size)

    # 3. 逐通道计算峰度
    channel_kurtosis = [
        kurtosis(reshaped_tensor[i]) for i in range(reshaped_tensor.shape[0])
    ]

    # 4. 计算通道峰度的平均值
    channel_kurtosis_avg = np.mean(channel_kurtosis)

    return global_kurtosis, channel_kurtosis, channel_kurtosis_avg

def calculate_kurtosis(tensor: torch.Tensor):
    tensor = tensor.squeeze(0)  # (sequence_length, hidden_size)

    # 计算每个通道 (hidden_size) 的均值和标准差
    mean = torch.mean(tensor, dim=1, keepdim=True)  # (1, hidden_size)
    std = torch.std(tensor, dim=1, keepdim=True) + 1e-8  # 避免除零错误

    # 标准化张量
    standardized_tensor = (tensor - mean) / std

    # 计算每个通道的峰度 (4 次方的均值)
    kurtosis = torch.mean(standardized_tensor ** 4, dim=1)  # (hidden_size,)

    return torch.mean(kurtosis).item()


def show_without_plt(mp, print_info=False):
    result = {}

    # print(mp.shape)
    all = mp.numel()
    # print(all)
    abs_mp = torch.abs(mp)

    bins = [0, 1e-10, 1e-09, 1e-08, 1e-07, 1e-06, 1e-05, 0.0001, 0.001, 0.01, 0.1, 1.0, 10.0, 100.0, 1000.0, 1000000.0]

    # 统计每个区间内的元素数量
    counts = []
    first = True
    for i in range(0,len(bins) - 1):
        lower_bound = bins[i]
        upper_bound = bins[i + 1]
        count = ((abs_mp >= lower_bound) & (abs_mp < upper_bound)).sum().item()
        if len(counts) == 0:
            counts.append(count)
            first = False
        elif count > 0 and first and counts[0] == 0:
            counts[0] = count
            first = False
            counts.append(0)
        else:
            counts.append(count)
    # print(counts)
    assert sum(counts) == all
    # print(sum(counts))


    # 8：1e-3
    # 7：1e-4


    low_value_counts_1e_3 = sum(counts[:8])
    low_value_counts_1e_4 = sum(counts[:7])

    counts_sum = sum(counts)

    lower_ratio_1e_3 = low_value_counts_1e_3 / counts_sum
    lower_ratio_1e_4 = low_value_counts_1e_4 / counts_sum

    # global_kurtosis, channel_kurtosis, channel_kurtosis_avg = compute_global_and_channel_kurtosis_numpy(mp)
    global_kurtosis, channel_kurtosis, channel_kurtosis_avg = 0,0 , calculate_kurtosis(mp)



    result["counts"] = counts
    result["bins"] = bins
    result["low_value_counts_1e_3"] = low_value_counts_1e_3
    result["low_value_counts_1e_4"] = low_value_counts_1e_4
    result["lower_ratio_1e_3"] = lower_ratio_1e_3
    result["lower_ratio_1e_4"] = lower_ratio_1e_4
    result["global_kurtosis"] = global_kurtosis
    result["channel_kurtosis_avg"] = channel_kurtosis_avg


    if print_info:
        print(result)

    return {
        "counts": counts,
        "bins": bins,
        "low_value_counts_1e_3": low_value_counts_1e_3,
        "low_value_counts_1e_4": low_value_counts_1e_4,
        "lower_ratio_1e_3": lower_ratio_1e_3,
        "lower_ratio_1e_4": lower_ratio_1e_4,
        "global_kurtosis": global_kurtosis,
        "channel_kurtosis_avg": channel_kurtosis_avg,
        # "channel_kurtosis": channel_kurtosis
    }

def get_hat():
    pass



class Evaler:
    def __init__(self, print_info=False,cal_kurtosis = False,cal_time=False):
        self.print_info = print_info
        self.cal_kurtosis = cal_kurtosis
        self.cal_time = cal_time

    def eval_tensor(self, tensor):
        result = {}

        # print(mp.shape)
        all = tensor.numel()
        # print(all)
        abs_mp = torch.abs(tensor)

        bins = [0, 1e-10, 1e-09, 1e-08, 1e-07, 1e-06, 1e-05, 0.0001, 0.001, 0.01, 0.1, 1.0, 10.0, 100.0, 1000.0,
                1000000.0]

        # 统计每个区间内的元素数量
        counts = []

        if self.cal_time:
            start_time = time.time()

        first = True
        for i in range(0, len(bins) - 1):
            lower_bound = bins[i]
            upper_bound = bins[i + 1]
            count = ((abs_mp >= lower_bound) & (abs_mp < upper_bound)).sum().item()
            if len(counts) == 0:
                counts.append(count)
                first = False
            elif count > 0 and first and counts[0] == 0:
                counts[0] = count
                first = False
                counts.append(0)
            else:
                counts.append(count)

        if self.cal_time:
            uesd_time = time.time() - start_time
            print(f"uesd_time: {uesd_time}")

        counts_sum = sum(counts)

        assert counts_sum == all

        # 8：1e-3
        # 7：1e-4

        low_value_counts_1e_3 = sum(counts[:8])
        low_value_counts_1e_4 = sum(counts[:7])

        lower_ratio_1e_3 = low_value_counts_1e_3 / counts_sum
        lower_ratio_1e_4 = low_value_counts_1e_4 / counts_sum

        if self.cal_kurtosis:
            global_kurtosis, channel_kurtosis, channel_kurtosis_avg = compute_global_and_channel_kurtosis_numpy(tensor)
            result["global_kurtosis"] = global_kurtosis
            result["channel_kurtosis_avg"] = channel_kurtosis_avg

        result["counts"] = counts
        result["bins"] = bins
        result["low_value_counts_1e_3"] = low_value_counts_1e_3
        result["low_value_counts_1e_4"] = low_value_counts_1e_4
        result["lower_ratio_1e_3"] = lower_ratio_1e_3
        result["lower_ratio_1e_4"] = lower_ratio_1e_4

        if self.print_info:
            print(result)

        return result

    def get_hat(self):
        result_hat = set()
        if self.cal_kurtosis:
            result_hat.add("global_kurtosis")
            result_hat.add("channel_kurtosis_avg")
        result_hat.add("counts")
        result_hat.add("bins")
        result_hat.add("low_value_counts_1e_3")
        result_hat.add("low_value_counts_1e_4")

        return result_hat

if __name__ == '__main__':
    # 创建一个随机张量
    fm = torch.cat([torch.load(f"ffn/second_line_input/tensor{i}.pth", weights_only=True) for i in range(32)]).to("cuda:0")
    rotation_matrix = torch.load("1-10-1 (1).pth", weights_only=True)
    rotated_fm = torch.matmul(fm, rotation_matrix)
    show_without_plt(fm, print_info=False)
    show_without_plt(rotated_fm, print_info=False)

    # load_tensors_1 = torch.load('original_input.pth', weights_only=True)
    # load_tensors_2 = torch.load('modified_input.pth', weights_only=True)

    # counts_1, counts_2 = [], []
    # bins = []
    # # 遍历加载的张量
    # for tensor1, tensor2 in zip(load_tensors_1, load_tensors_2):
    #     temp1 = show_without_plt(tensor1)
    #     temp2 = show_without_plt(tensor2)
    #     counts_1.append(temp1["counts"])
    #     counts_2.append(temp2["counts"])
    #     bins = temp1["bins"]  # 假设两组数据的 bins 是一致的



